Skip to content

Fix incorrect pairing in sample_plan_with_scipy#178

Merged
atong01 merged 1 commit into
atong01:mainfrom
zsteve:patch-1
Apr 20, 2026
Merged

Fix incorrect pairing in sample_plan_with_scipy#178
atong01 merged 1 commit into
atong01:mainfrom
zsteve:patch-1

Conversation

@zsteve
Copy link
Copy Markdown
Contributor

@zsteve zsteve commented Apr 20, 2026

What does this PR do?

Currently sample_plan_with_scipy returns incorrectly paired points due to an indexing bug. Should be

        pi_x0 = x0
        pi_x1 = x1[j]

not

       pi_x0 = x0[j]
       pi_x1 = x1

Before submitting

  • Did you make sure title is self-explanatory and the description concisely explains the PR?
  • Did you make sure your PR does only one thing, instead of bundling different changes together?
  • Did you list all the breaking changes introduced by this pull request?
  • Did you test your PR locally with pytest command?
  • Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃

Summary by Sourcery

Bug Fixes:

  • Correct point pairing in sample_plan_with_scipy by applying the assignment indices to the target points instead of the source points.

@sourcery-ai
Copy link
Copy Markdown
Contributor

sourcery-ai Bot commented Apr 20, 2026

Reviewer's guide (collapsed on small PRs)

Reviewer's Guide

Fixes the point pairing logic in sample_plan_with_scipy by correcting which tensor is indexed by the assignment returned from SciPy’s linear_sum_assignment, so that x1 is permuted instead of x0.

Sequence diagram for corrected sample_plan_with_scipy pairing

sequenceDiagram
    participant Caller
    participant SamplePlan
    participant ScipyOptimize

    Caller->>SamplePlan: sample_plan_with_scipy(x0, x1)
    SamplePlan->>SamplePlan: compute cost matrix M from x0, x1
    SamplePlan->>SamplePlan: optionally normalize M
    SamplePlan->>ScipyOptimize: linear_sum_assignment(M)
    ScipyOptimize-->>SamplePlan: j
    SamplePlan->>SamplePlan: pi_x0 = x0
    SamplePlan->>SamplePlan: pi_x1 = x1[j]
    SamplePlan-->>Caller: (pi_x0, pi_x1)
Loading

File-Level Changes

Change Details Files
Correct the indexing in sample_plan_with_scipy so that the optimal assignment indices are applied to x1 instead of x0.
  • Keep x0 unchanged when forming the transport plan output
  • Index x1 with the assignment indices j returned from linear_sum_assignment to produce correctly paired points
  • Return the corrected (pi_x0, pi_x1) pair from sample_plan_with_scipy
torchcfm/optimal_transport.py

Tips and commands

Interacting with Sourcery

  • Trigger a new review: Comment @sourcery-ai review on the pull request.
  • Continue discussions: Reply directly to Sourcery's review comments.
  • Generate a GitHub issue from a review comment: Ask Sourcery to create an
    issue from a review comment by replying to it. You can also reply to a
    review comment with @sourcery-ai issue to create an issue from it.
  • Generate a pull request title: Write @sourcery-ai anywhere in the pull
    request title to generate a title at any time. You can also comment
    @sourcery-ai title on the pull request to (re-)generate the title at any time.
  • Generate a pull request summary: Write @sourcery-ai summary anywhere in
    the pull request body to generate a PR summary at any time exactly where you
    want it. You can also comment @sourcery-ai summary on the pull request to
    (re-)generate the summary at any time.
  • Generate reviewer's guide: Comment @sourcery-ai guide on the pull
    request to (re-)generate the reviewer's guide at any time.
  • Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
    pull request to resolve all Sourcery comments. Useful if you've already
    addressed all the comments and don't want to see them anymore.
  • Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
    request to dismiss all existing Sourcery reviews. Especially useful if you
    want to start fresh with a new review - don't forget to comment
    @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

  • Enable or disable review features such as the Sourcery-generated pull request
    summary, the reviewer's guide, and others.
  • Change the review language.
  • Add, remove or edit custom review instructions.
  • Adjust other review settings.

Getting Help

Copy link
Copy Markdown
Contributor

@sourcery-ai sourcery-ai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey - I've reviewed your changes and they look great!


Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

@atong01
Copy link
Copy Markdown
Owner

atong01 commented Apr 20, 2026

Good catch! Thanks!

@atong01 atong01 merged commit cacd4dc into atong01:main Apr 20, 2026
19 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants